# 定义一个向量数据库的类，实现数据库初始化，简表，增删改查等基本操作
# 首先导入一些需要用到的库
from pymilvus import (
    connections,
    utility,
    FieldSchema, CollectionSchema, DataType,
    Collection,
)


class MilvusDatabase:

    # 一个数据库初始化的函数
    def __init__(self, dim, host='localhost', port=19530, user='root', password='', database="default"):
        # 与数据库建立连接
        connections.connect(db_name=database, host=host, port=port, user=user, password=password)
        # 如果不存在对应的库就进行创建
        if not self.is_collection_exist():
            # 创建collection
            # 创建的是集合，集合里面有领域（就是相当于列）

            fields = [
                FieldSchema(name='id', dtype=DataType.INT64, is_primary=True, auto_id=True),
                FieldSchema(name='file_name', dtype=DataType.VARCHAR, max_length=200),
                FieldSchema(name='page', dtype=DataType.INT64),
                FieldSchema(name='content_code', dtype=DataType.VARCHAR, max_length=250),
                FieldSchema(name='embeddings', dtype=DataType.FLOAT_VECTOR, dim=dim),  # 向量有一个维度的参数
                FieldSchema(name='content', dtype=DataType.VARCHAR, max_length=65535),
            ]

            # 创建这个集合
            schema = CollectionSchema(fields, description="content_embedding_db")
            self.fileEmbeddingDB = Collection("content_embedding_db", schema, consistency_level="Strong")

            # 在集合上创建索引
            index = {
                "index_type": "IVF_FLAT",  # 索引的类型
                "metric_type": "L2",  # 指标的类型
                "params": {"nlist": 128},
            }
            self.fileEmbeddingDB.create_index("embeddings", index)

            self.fileEmbeddingDB.load()

        else:
            self.fileEmbeddingDB = Collection("content_embedding_db")
            self.fileEmbeddingDB.load()

    # 一个向数据库中插入数据的函数
    def insert_data(self, entities, partition_name=None):
        if partition_name is not None:
            if not self.fileEmbeddingDB.has_partition(partition_name):
                # 如果不存在这个分区的话就创建一下
                self.fileEmbeddingDB.create_partition(partition_name)
        insert_result = self.fileEmbeddingDB.insert(entities, partition_name=partition_name)
        self.fileEmbeddingDB.flush()

    # 一个从数据库中删除数据的函数
    def delete_data(self, expr, partition_name=None):
        delete_result = self.fileEmbeddingDB.delete(expr, partition_name=partition_name)
        self.fileEmbeddingDB.flush()

    # 基于向量相似度对向量进行查询
    def search_embedding(self, vectors_to_search, limit, expr=None, output_fields=None, partition_names=None):
        """
        vectors_to_search是指的要从数据库里查相似度的向量
        limit是指的每次返回几条最相似的数据
        output_files指的是要返回的哪些列组成的列表
        """
        search_params = {
            "metric_type": "L2",
            "params": {"nprobe": 10},
        }
        result = self.fileEmbeddingDB.search(vectors_to_search, "embeddings", search_params, expr=expr, limit=limit,
                                             output_fields=output_fields, partition_names=partition_names)

        # 得到查询的返回的值
        res = [[([hit.entity.get(filed) for filed in output_fields], hit.distance) for hit in hits] for hits in result]
        return res

    # 基于标量进行查询（只查文件名或者是只查路径，或者是基于文件名的相似查询）
    def search_scaler(self, expr, output_fields=None, partition_names=None, limit=None):
        if limit is None:
            limit = 20
        # expr举例："book_id in [2,4,6,8]",><表达式也能用，还有字符串条件模糊查询时可以设置"column_name.str.contains("abc")"==找到包含特定字符串的那些行
        result = self.fileEmbeddingDB.query(expr=expr, output_fields=output_fields,
                                            partition_names=partition_names, limit=limit)
        try:
            res = [[hits[filed] for filed in output_fields] for hits in result]
        except:
            res = []
        return res

    # 直接进行数据插入
    def insert_embedding(self, file_name, page, content_code, new_embedding, content, partition_name=None):
        # 直接向集合中插入数据
        entities = [
            [file_name],
            [page],
            [content_code],
            [new_embedding],
            [content]
        ]
        self.insert_data(entities, partition_name=partition_name)

    def insert_embedding_batch(self, data=None, partition_name=None):
        if data is None:
            data = []
        if partition_name is not None:
            if not self.fileEmbeddingDB.has_partition(partition_name):
                # 如果不存在这个分区的话就创建一下
                self.fileEmbeddingDB.create_partition(partition_name)
        insert_result = self.fileEmbeddingDB.insert(data, partition_name=partition_name)
        self.fileEmbeddingDB.flush()

    # 删除集合
    @staticmethod
    def drop_collection():
        utility.drop_collection('content_embedding_db')

    # 查看集合中实体的数量
    def get_entity_num(self):
        return self.fileEmbeddingDB.num_entities

    # 查看集合是否存在
    @staticmethod
    def is_collection_exist():
        return utility.has_collection("content_embedding_db")

    # 列出所有的集合名称
    @staticmethod
    def show_all_collection_name():
        return utility.list_collections()
